import torch
import matplotlib.pyplot as plt
from torch import nn,optim
### 如果想要添加到GPU上进行运行
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


##### 准备数据
x = torch.rand([50,1])
# x = torch.rand([50,1]).to(device)
y = 10 * x + 8

##### 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel,self).__init__()
        self.linear = nn.Linear(1,1)
    def forward(self,x):
        out = self.linear(x)
        return out


##### 实例化
model = MyModel()
# model = MyModel().to(device)
optimizer = optim.SGD(model.parameters(),lr=0.01)
criterion = nn.MSELoss()

##### 进行梯度下降并更新
for i in range(5000):
    y_pred = model(x)
    loss = criterion(y,y_pred)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if i%500 ==0:
        print('epoch:{},loss:{:.7f}'.format(i,loss))

##### 模型评估
model.eval()
y_pred = model(x)
plt.scatter(x.data,y.data,c = 'b')
plt.plot(x.data,y_pred.data,c = 'r')
plt.title('Result pic')
plt.show()


